import json
import argparse


TEMPLATES = [
    "a photo of {crop} leaves diseased by {disease}",
    "{crop} leaves showing symptoms of {disease}",
    "leaf symptoms consistent with {disease} on {crop}",
    "healthy {crop} leaves with normal color",
]


def heuristics_for_class(cname: str):
    # crude split into crop and disease if possible
    # users can edit output to improve
    crop = "plant"
    disease = cname
    # try split patterns
    parts = cname.split()
    if len(parts) >= 2:
        # heuristic: take last token(s) as disease
        crop = parts[0]
        disease = cname
    prompts = []
    for t in TEMPLATES:
        prompts.append(t.format(crop=crop, disease=disease))
    return prompts


def main():
    ap = argparse.ArgumentParser(description='Build prompts.json from classes.txt')
    ap.add_argument('--classes', required=True, help='Path to classes.txt')
    ap.add_argument('--out', required=True, help='Output prompts.json')
    args = ap.parse_args()

    with open(args.classes, 'r', encoding='utf-8') as f:
        classes = [line.strip() for line in f if line.strip()]

    prompts = {}
    for cname in classes:
        prompts[cname] = heuristics_for_class(cname)

    with open(args.out, 'w', encoding='utf-8') as f:
        json.dump(prompts, f, ensure_ascii=False, indent=2)
    print(f"Saved prompts.json to {args.out} (classes={len(classes)})")


if __name__ == '__main__':
    main()